import decimal
from Module import Resnet101 as Net # 导入模型
# from torch.nn import CrossEntropyLoss as Loss # 导入损失函数
from torch.nn import BCEWithLogitsLoss as Loss # 导入损失函数
from torch.optim import AdamW as Optim # 导入优化器
from ImgDataset import ImgDataset, GetData # 导入数据生成器
from EarlyStopping import EarlyStopping # 导入早停器
from torch import autocast # 导入半精度控制器
from torch.utils.data import DataLoader # 导入数据生成器
from pathlib import Path
import torch
import numpy as np
from Plot import PlotData

DEVICE =  "cuda" if torch.cuda.is_available() else "cpu" #b 定义设备常量
IMG_SIZE = 256
EPOCHS = 200
np.set_printoptions(linewidth=np.inf) # type: ignore # 设置np数组打印的宽度


# 创建数据
p = r'/home/qhj_work/dataset/indoorCVPR_09/Images' # 数据集地址
p1 = r'/mnt/e/VS Project/car parts 50/car parts 50'
dataset = GetData(p1, valid_test_size=5, seed=12) # 构建训练集，验证集
train_dataset = ImgDataset(dataset.train_dataset,
                           label_set=dataset.lebel_onehot,
                           mode='train',
                           img_size=IMG_SIZE,
                           img_rise=True,
                           shuffle=True,
                           seed=3) # 创建训练集
valid_dataset = ImgDataset(dataset.valid_dataset,
                        label_set=dataset.lebel_onehot,
                        mode='valid',
                        img_size=IMG_SIZE,
                        shuffle=True,
                        seed=4) # 创建验证集


# 模型实例化
module = Net(num_classes=len(dataset.lebel_onehot))


# 定义损失函数
loss = Loss()

# 定义优化器
optim = Optim(params=module.parameters(),
              lr=0.0001,
              weight_decay=0.05
              )


# 定义数据集生成器
train_loader = DataLoader(train_dataset, batch_size=16,
                          num_workers=8, # 调用进程数，
                          # linux系统中可以使用多个子进程加载数据，而在windows系统中不能。在windows系统中需要将进程数设置为单进程0
                          )
valid_loader = DataLoader(valid_dataset, batch_size=16,
                          num_workers=8, # 调用进程数
                          )


# 定义早停器
earlystop_loss = EarlyStopping(patience=20, autodelta=True, recently=True)


# 定义训练
def train(epoch):
    train_loss = 0
    gt_labels = []
    pred_labels = []
    module.train() # 将模型设置为训练状态
    for k, (data, label) in enumerate(train_loader):
        optim.zero_grad()
        with autocast(DEVICE, torch.bfloat16): # 采用半精度前向传播
            if torch.cuda.is_available():
                data, label = data.cuda(), label.cuda() # 迁移gpu
                module.cuda()
            t = module(data.to(torch.float32))
            l = loss(t, label.to(torch.float32)) * 100
        l.backward()
        optim.step()
        train_loss += l.item()*data.size(0)
        preds = torch.argmax(t, dim=1) # 返回最大值的索引
        gt_labels.append(label.cpu().numpy())
        pred_labels.append(preds.cpu().numpy())
        if k % 50 == 0:
            print('Epoch:{}, Train batch:{}, loss_mse:{}'.format(
                epoch, k, l.data
            ))
    else:
        index = np.random.choice([i for i in range(label.size(0))])
        show_t = t[index].cpu().detach().to(torch.float).numpy().round(decimals=2)
        t_max = np.argmax(show_t)
        show_l = label[index].cpu().detach().to(torch.float).numpy().round(decimals=2)
        l_max = np.argmax(show_l)
        print('Train pred_max_index:{}, pred:\n{}\nValid label_max_index:{}, label:\n{}'.format(
            t_max, show_t, l_max, show_l
            ))
    train_loss = train_loss/len(train_loader.dataset) # type: ignore
    gt_labels, pred_labels = np.concatenate(gt_labels), np.concatenate(pred_labels)
    acc = np.sum(gt_labels[pred_labels == np.argmax(gt_labels, axis=1)])/len(pred_labels) * 100 # 计算准确率
    print('Epoch: {} \tTraining Loss: {:.6f}\tTraining acc: {:.6f}'.format(epoch, train_loss, acc))
    return train_loss, acc

# 定义验证
def valid(epoch):
    module.eval()
    valid_loss = 0
    gt_labels = []
    pred_labels = []
    with torch.no_grad():
        for k, (data, label) in enumerate(valid_loader):
            with autocast(DEVICE, torch.bfloat16): # 采用半精度前向传播
                if torch.cuda.is_available():
                    data, label = data.cuda(), label.cuda() # 迁移gpu
                    module.cuda()
                t = module(data.to(torch.float32))
                l = loss(t, label.to(torch.float32)) * 100
            preds = torch.argmax(t, dim=1) # 返回最大值的索引
            gt_labels.append(label.cpu().numpy())
            pred_labels.append(preds.cpu().numpy())
            valid_loss += l.item()*data.size(0)
        else:
            index = np.random.choice([i for i in range(label.size(0))])
            show_t = t[index].cpu().detach().to(torch.float).numpy().round(decimals=2)
            t_max = np.argmax(show_t)
            show_l = label[index].cpu().detach().to(torch.float).numpy().round(decimals=2)
            l_max = np.argmax(show_l)
            print('Valid pred_max_index:{}, pred:\n{}\nValid label_max_index:{}, label:\n{}'.format(
                t_max, show_t, l_max, show_l
                ))
        valid_loss = valid_loss/len(valid_loader.dataset) # type: ignore
        gt_labels, pred_labels = np.concatenate(gt_labels), np.concatenate(pred_labels)
        acc = np.sum(gt_labels[pred_labels == np.argmax(gt_labels, axis=1)])/len(pred_labels) * 100 # 计算准确率
        # 早停器判断
        if epoch >= EPOCHS / 4 or epoch > 30:
            # earlystop_loss(valid_loss)
            pass
        print('Epoch: {} \tValidation Loss: {:.6f}\tAccuracy: {:6f}'.format(epoch, valid_loss, acc))
    return valid_loss, acc



epochs = [i for i in range(1, EPOCHS+1)]
e_list = []
train_loss, valid_loss, acc, acc_t = [], [], [], []
for epoch in epochs:
    e_list.append(epoch)
    t, t_a = train(epoch)
    v, a = valid(epoch)
    train_loss.append(t)
    valid_loss.append(v)
    acc.append(a)
    acc_t.append(t_a)
    if earlystop_loss.early_stop:
        print("Training Stop, cause the error loss has continued to rise")
        break

# 画图
showdata = PlotData(e_list, y=[train_loss, valid_loss, acc, acc_t], fig_name='training error', legend=['train loss', 'valid loss', 'valid acc', 'train acc'])
showdata()

# 保存模型
while True:
    q = input('请确认是否保存模型[y]/[n]:')
    if q == 'y':
        current_dir = Path.cwd()
        m_dir = current_dir / 'model.pth'
        state_dir = current_dir / 'model_weights.pth'
        torch.save(module, m_dir)
        torch.save(module.state_dict(), state_dir)
        break
    elif q== 'n':
        break